# Applies one local Windows policy value using simple FileWave/customer environment variables.
# Optional: place LGPO.exe beside this script to import the generated LGPO text into local policy.

$ErrorActionPreference = "Stop"

function Get-ScriptRoot {
    if ($PSScriptRoot) {
        return $PSScriptRoot
    }

    return Split-Path -Parent $MyInvocation.MyCommand.Path
}

function Get-Setting {
    param(
        [string[]]$Names,
        [AllowNull()][string]$Default = $null
    )

    foreach ($Name in $Names) {
        foreach ($Target in @("Process", "Machine", "User")) {
            $Value = [Environment]::GetEnvironmentVariable($Name, $Target)
            if ($null -ne $Value) {
                return $Value
            }
        }
    }

    return $Default
}

function Get-BooleanSetting {
    param(
        [string[]]$Names,
        [bool]$Default
    )

    $Value = Get-Setting -Names $Names
    if ([string]::IsNullOrWhiteSpace($Value)) {
        return $Default
    }

    switch ($Value.Trim().ToLowerInvariant()) {
        { $_ -in @("1", "true", "yes", "y", "on") } { return $true }
        { $_ -in @("0", "false", "no", "n", "off") } { return $false }
        default { throw "Invalid true/false value '$Value' for $($Names[0])." }
    }
}

function Resolve-PolicyLocation {
    param(
        [string]$Scope,
        [string]$Key
    )

    if ([string]::IsNullOrWhiteSpace($Key)) {
        throw "FW_POLICY_KEY is required."
    }

    $ResolvedScope = if ([string]::IsNullOrWhiteSpace($Scope)) { "Computer" } else { $Scope.Trim() }
    $ResolvedKey = $Key.Trim()

    if ($ResolvedKey -match "^(HKLM:?\\|HKEY_LOCAL_MACHINE\\)") {
        $ResolvedScope = "Computer"
        $ResolvedKey = $ResolvedKey -replace "^(HKLM:?\\|HKEY_LOCAL_MACHINE\\)", ""
    }
    elseif ($ResolvedKey -match "^(HKCU:?\\|HKEY_CURRENT_USER\\)") {
        $ResolvedScope = "User"
        $ResolvedKey = $ResolvedKey -replace "^(HKCU:?\\|HKEY_CURRENT_USER\\)", ""
    }

    switch ($ResolvedScope.Trim().ToLowerInvariant()) {
        { $_ -in @("computer", "machine", "hklm") } { $ResolvedScope = "Computer" }
        { $_ -in @("user", "hkcu") } { $ResolvedScope = "User" }
        default { throw "FW_POLICY_SCOPE must be Computer or User." }
    }

    return @{
        Scope = $ResolvedScope
        Key = $ResolvedKey.TrimStart("\")
    }
}

function Resolve-PolicyType {
    param([string]$Type)

    if ([string]::IsNullOrWhiteSpace($Type)) {
        throw "FW_POLICY_TYPE is required when applying a policy."
    }

    switch ($Type.Trim().ToUpperInvariant().Replace("-", "_")) {
        "DWORD" { return @{ RegistryType = "DWord"; LGPOPrefix = "DWORD" } }
        "REG_DWORD" { return @{ RegistryType = "DWord"; LGPOPrefix = "DWORD" } }
        "QWORD" { return @{ RegistryType = "QWord"; LGPOPrefix = "QWORD" } }
        "REG_QWORD" { return @{ RegistryType = "QWord"; LGPOPrefix = "QWORD" } }
        "SZ" { return @{ RegistryType = "String"; LGPOPrefix = "SZ" } }
        "STRING" { return @{ RegistryType = "String"; LGPOPrefix = "SZ" } }
        "REG_SZ" { return @{ RegistryType = "String"; LGPOPrefix = "SZ" } }
        "EXPAND_SZ" { return @{ RegistryType = "ExpandString"; LGPOPrefix = "EXSZ" } }
        "REG_EXPAND_SZ" { return @{ RegistryType = "ExpandString"; LGPOPrefix = "EXSZ" } }
        "MULTI_SZ" { return @{ RegistryType = "MultiString"; LGPOPrefix = "MULTISZ" } }
        "REG_MULTI_SZ" { return @{ RegistryType = "MultiString"; LGPOPrefix = "MULTISZ" } }
        default { throw "FW_POLICY_TYPE '$Type' is not supported. Use DWORD, QWORD, SZ, EXPAND_SZ, or MULTI_SZ." }
    }
}

function ConvertTo-RegistryValue {
    param(
        [string]$RegistryType,
        [string]$Data
    )

    switch ($RegistryType) {
        "DWord" {
            if ($Data -match "^0x") {
                return [Convert]::ToUInt32($Data.Substring(2), 16)
            }

            return [UInt32]$Data
        }
        "QWord" {
            if ($Data -match "^0x") {
                return [Convert]::ToUInt64($Data.Substring(2), 16)
            }

            return [UInt64]$Data
        }
        "String" { return [string]$Data }
        "ExpandString" { return [string]$Data }
        "MultiString" { return [string[]]($Data -split "\|") }
        default { throw "Direct registry enforcement does not support '$RegistryType'." }
    }
}

function Get-RegistryPath {
    param(
        [string]$Scope,
        [string]$Key
    )

    if ($Scope -eq "Computer") {
        return "HKLM:\$Key"
    }

    return "HKCU:\$Key"
}

$ScriptRoot = Get-ScriptRoot
$Action = (Get-Setting -Names @("FW_POLICY_ACTION") -Default "SET").Trim().ToUpperInvariant()
$Location = Resolve-PolicyLocation `
    -Scope (Get-Setting -Names @("FW_POLICY_SCOPE") -Default "Computer") `
    -Key (Get-Setting -Names @("FW_POLICY_KEY"))

$ValueName = Get-Setting -Names @("FW_POLICY_VALUE")
if ([string]::IsNullOrWhiteSpace($ValueName)) {
    throw "FW_POLICY_VALUE is required."
}

if ($Action -notin @("SET", "DELETE")) {
    throw "FW_POLICY_ACTION must be SET or DELETE."
}

$LGPO = Get-Setting -Names @("FW_LGPO_EXE") -Default (Join-Path $ScriptRoot "LGPO.exe")
$LGPOTextPath = Get-Setting -Names @("FW_LGPO_TEXT_PATH") -Default (Join-Path $ScriptRoot "GeneratedPolicy.txt")
$RunGPUpdate = Get-BooleanSetting -Names @("FW_RUN_GPUPDATE") -Default $true
$DirectRegistryFallback = Get-BooleanSetting -Names @("FW_DIRECT_REGISTRY_FALLBACK") -Default $true

if ($Action -eq "DELETE") {
    $DataLine = "DELETE"
    $PolicyType = $null
    $Data = $null
}
else {
    $Data = Get-Setting -Names @("FW_POLICY_DATA")
    if ($null -eq $Data) {
        throw "FW_POLICY_DATA is required when applying a policy."
    }

    $PolicyType = Resolve-PolicyType -Type (Get-Setting -Names @("FW_POLICY_TYPE"))
    $DataLine = "$($PolicyType.LGPOPrefix):$Data"
}

$LGPOText = @(
    $Location.Scope
    $Location.Key
    $ValueName
    $DataLine
) -join "`r`n"

Set-Content -Path $LGPOTextPath -Value $LGPOText -Encoding ASCII
Write-Host "Generated LGPO text policy: $LGPOTextPath"

if (Test-Path $LGPO -PathType Leaf) {
    Write-Host "Applying policy with LGPO.exe"
    & $LGPO /t $LGPOTextPath /v

    if ($LASTEXITCODE -ne 0) {
        throw "LGPO.exe failed with exit code $LASTEXITCODE"
    }
}
else {
    Write-Host "LGPO.exe was not found at $LGPO."
}

if ($DirectRegistryFallback) {
    if ($Location.Scope -ne "Computer") {
        Write-Host "Skipping direct registry fallback for User policy. Use LGPO.exe for user policies in FileWave."
    }
    elseif ($Action -eq "DELETE") {
        $RegistryPath = Get-RegistryPath -Scope $Location.Scope -Key $Location.Key
        Write-Host "Removing registry policy $RegistryPath\$ValueName"
        Remove-ItemProperty -Path $RegistryPath -Name $ValueName -ErrorAction SilentlyContinue
    }
    else {
        $RegistryPath = Get-RegistryPath -Scope $Location.Scope -Key $Location.Key
        Write-Host "Writing registry policy $RegistryPath\$ValueName"
        New-Item -Path $RegistryPath -Force | Out-Null
        New-ItemProperty `
            -Path $RegistryPath `
            -Name $ValueName `
            -PropertyType $PolicyType.RegistryType `
            -Value (ConvertTo-RegistryValue -RegistryType $PolicyType.RegistryType -Data $Data) `
            -Force | Out-Null
    }
}

if ($RunGPUpdate) {
    gpupdate.exe /force
}

Write-Host "Policy action completed: $Action $($Location.Scope)\$($Location.Key)\$ValueName"
